# coding=utf-8
# Copyright (C) yhy team - All Rights Reserved
#
# @Version:   3.10.4
# @Software:  PyCharm
# @FileName:  可视化激活函数的学习效果.py
# @CTime:     2022/12/8 11:56   
# @Author:    yhy
# @Email:     yhy@yhy.com
# @UTime:     2022/12/8 11:56
#
# @Description:
#     
#     xxx
#
import codecs
import logging
from typing import List, Dict, Optional
import json
import math
import os
import urllib.request
import warnings
from urllib.error import HTTPError
import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import seaborn as sns
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torch.utils.data as data
import torchvision
from torchvision import transforms
from torchvision.datasets import FashionMNIST
from tqdm import tqdm

matplotlib.use('TkAgg')

logger = logging.getLogger(__name__)


def set_seed(seed):
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():  # GPU operation have separate seed
        torch.cuda.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)
        torch.backends.cudnn.determinstic = True
        torch.backends.cudnn.benchmark = False


class ActivationFunction(nn.Module):
    def __init__(self):
        super().__init__()
        self.name = self.__class__.__name__
        self.config = {"name": self.name}


class Sigmoid(ActivationFunction):
    def forward(self, x):
        return 1 / (1 + torch.exp(-x))


class Tanh(ActivationFunction):
    def forward(self, x):
        x_exp, neg_x_exp = torch.exp(x), torch.exp(-x)
        return (x_exp - neg_x_exp) / (x_exp + neg_x_exp)


class ReLU(ActivationFunction):
    def forward(self, x):
        return x * (x > 0).float()


class LeakyReLU(ActivationFunction):
    def __init__(self, alpha=0.1):
        super().__init__()
        self.config["alpha"] = alpha

    def forward(self, x):
        return torch.where(x > 0, x, self.config["alpha"] * x)


class ELU(ActivationFunction):
    def __init__(self, alpha=1):
        super().__init__()
        self.config["alpha"] = alpha

    def forward(self, x):
        # elu(x) = max(0, x) + min(0, alpha * ( exp(x)-1 )
        return torch.where(x > 0, x, self.config['alpha'] * (torch.exp(x) - 1))


class SiLU(ActivationFunction):
    def forward(self, x):
        return x * torch.sigmoid(x)


class SELU(ActivationFunction):
    def forward(self, x):
        scale = 1.0507009873554804934193349852946
        alpha = 1.6732632423543772848170429916717
        return scale * (torch.where(x > 0, x, alpha * (torch.exp(x) - 1)))


class GELU(ActivationFunction):
    def forward(self, x):
        return 0.5 * x * (1.0 + torch.erf(x / math.sqrt(2.0)))


class NewGELU(ActivationFunction):
    def forward(self, x):
        return 0.5 * x * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi) * (x + 0.044715 * torch.pow(x, 3.0))))


class FastGELU(ActivationFunction):
    def forward(self, x):
        return 0.5 * x * (1.0 + torch.tanh(x * 0.7978845608 * (1.0 + 0.044715 * x * x)))


class QuickGELU(ActivationFunction):
    def forward(self, x):
        return x * torch.sigmoid(1.702 * x)


class Mish(ActivationFunction):
    def forward(self, x):
        return x * torch.tanh(F.softplus(x))


def get_grads(act_fn, x):
    """Computes the gradients of an activation function at specified positions.

    Args:
        act_fn: An object of the class "ActivationFunction" with an implemented forward pass.
        x: 1D input tensor.
    Returns:
        A tensor with the same size of x containing the gradients of act_fn at x.
    """
    x = x.clone().requires_grad_()  # Mark the input as tensor for which we want to store gradients
    out = act_fn(x)
    out.sum().backward()  # Summing results in an equal gradient flow to each element in x
    return x.grad  # Accessing the gradients of x by "x.grad"


def vis_act_fn(act_fn, ax, x):
    # Run activation function
    y = act_fn(x)
    y_grads = get_grads(act_fn, x)
    # Push x, y and gradients back to cpu for plotting
    x, y, y_grads = x.cpu().numpy(), y.cpu().numpy(), y_grads.cpu().numpy()
    # Plotting
    ax.plot(x, y, linewidth=2, label="ActFn")
    ax.plot(x, y_grads, linewidth=2, label="Gradient")
    ax.set_title(act_fn.name)
    ax.legend()
    ax.set_ylim(-1.5, x.max())

class BaseNetwork(nn.Module):
    def __init__(self, act_fn, input_size=784, num_classes=10, hidden_sizes=[512, 256, 256, 128]):
        """
        Args:
            act_fn: Object of the activation function that should be used as non-linearity in the network.
            input_size: Size of the input images in pixels
            num_classes: Number of classes we want to predict
            hidden_sizes: A list of integers specifying the hidden layer sizes in the NN
        """
        super().__init__()

        # Create the network based on the specified hidden sizes
        layers = []
        layer_sizes = [input_size] + hidden_sizes
        layer_size_last = layer_sizes[0]
        for layer_size in layer_sizes[1:]:
            layers += [nn.Linear(layer_size_last, layer_size), act_fn]
            layer_size_last = layer_size
        layers += [nn.Linear(layer_sizes[-1], num_classes)]
        # nn.Sequential summarizes a list of modules into a single module, applying them in sequence
        self.layers = nn.Sequential(*layers)

        # We store all hyperparameters in a dictionary for saving and loading of the model
        self.config = {
            "act_fn": act_fn.config,
            "input_size": input_size,
            "num_classes": num_classes,
            "hidden_sizes": hidden_sizes,
        }

    def forward(self, x):
        x = x.view(x.size(0), -1)  # Reshape images to a flat vector
        out = self.layers(x)
        return out


def plot_all_actfn(act_fn_by_name):
    # Add activation functions if wanted
    act_fns = [act_fn() for act_fn in act_fn_by_name.values()]
    x = torch.linspace(-5, 5, 1000)  # Range on which we want to visualize the activation functions
    # Plotting
    cols = 4
    rows = math.ceil(len(act_fns) / float(cols))
    fig, ax = plt.subplots(rows, cols, figsize=(cols * 4, rows * 4))
    for i, act_fn in enumerate(act_fns):
        vis_act_fn(act_fn, ax[divmod(i, cols)], x)
    fig.subplots_adjust(hspace=0.3)
    plt.show()


def plot_fashionmnist_example():
    exmp_imgs = [train_set[i][0] for i in range(16)]
    # Organize the images into a grid for nicer visualization
    img_grid = torchvision.utils.make_grid(torch.stack(exmp_imgs, dim=0), nrow=4, normalize=True, pad_value=0.5)
    img_grid = img_grid.permute(1, 2, 0)

    plt.figure(figsize=(8, 8))
    plt.title("FashionMNIST examples")
    plt.imshow(img_grid)
    plt.axis("off")
    plt.show()
    plt.close()



def visualize_gradients(net, color="C0"):
    """
    Args:
        net: Object of class BaseNetwork
        color: Color in which we want to visualize the histogram (for easier separation of activation functions)
    """
    net.eval()
    small_loader = data.DataLoader(train_set, batch_size=256, shuffle=False)
    imgs, labels = next(iter(small_loader))
    imgs, labels = imgs.to(device), labels.to(device)

    # Pass one batch through the network, and calculate the gradients for the weights
    net.zero_grad()
    preds = net(imgs)
    loss = F.cross_entropy(preds, labels)
    loss.backward()
    # We limit our visualization to the weight parameters and exclude the bias to reduce the number of plots
    grads = {
        name: params.grad.data.view(-1).cpu().clone().numpy()
        for name, params in net.named_parameters()
        if "weight" in name
    }
    net.zero_grad()

    # Plotting
    columns = len(grads)
    fig, ax = plt.subplots(1, columns, figsize=(columns * 3.5, 2.5))
    fig_index = 0
    for key in grads:
        key_ax = ax[fig_index % columns]
        sns.histplot(data=grads[key], bins=30, ax=key_ax, color=color, kde=True)
        key_ax.set_title(str(key))
        key_ax.set_xlabel("Grad magnitude")
        fig_index += 1
    fig.suptitle(
        f"Gradient magnitude distribution for activation function {net.config['act_fn']['name']}", fontsize=14, y=1.05
    )
    fig.subplots_adjust(wspace=0.45)
    plt.show()
    plt.close()




def train_model(net, model_name, max_epochs=10, patience=7, batch_size=256, overwrite=False):
    """Train a model on the training set of FashionMNIST.

    Args:
        net: Object of BaseNetwork
        model_name: (str) Name of the model, used for creating the checkpoint names
        max_epochs: Number of epochs we want to (maximally) train for
        patience: If the performance on the validation set has not improved for #patience epochs, we stop training early
        batch_size: Size of batches used in training
        overwrite: Determines how to handle the case when there already exists a checkpoint. If True, it will be overwritten. Otherwise, we skip training.
    """
    file_exists = False
    if file_exists and not overwrite:
        print("Model file already exists. Skipping training...")
    else:
        if file_exists:
            print("Model file exists, but will be overwritten...")

        # Defining optimizer, loss and data loader
        optimizer = optim.SGD(net.parameters(), lr=1e-2, momentum=0.9)  # Default parameters, feel free to change
        loss_module = nn.CrossEntropyLoss()
        train_loader_local = data.DataLoader(
            train_set, batch_size=batch_size, shuffle=True, drop_last=True, pin_memory=True
        )

        val_scores = []
        best_val_epoch = -1
        for epoch in range(max_epochs):
            ############
            # Training #
            ############
            net.train()
            true_preds, count = 0.0, 0
            for imgs, labels in tqdm(train_loader_local, desc=f"Epoch {epoch+1}", leave=False):
                imgs, labels = imgs.to(device), labels.to(device)  # To GPU
                optimizer.zero_grad()  # Zero-grad can be placed anywhere before "loss.backward()"
                preds = net(imgs)
                loss = loss_module(preds, labels)
                loss.backward()
                optimizer.step()
                # Record statistics during training
                true_preds += (preds.argmax(dim=-1) == labels).sum()
                count += labels.shape[0]
            train_acc = true_preds / count

            ##############
            # Validation #
            ##############
            val_acc = test_model(net, val_loader)
            val_scores.append(val_acc)
            print(
                f"[Epoch {epoch+1}] Training accuracy: {train_acc*100.0:05.2f}%, Validation accuracy: {val_acc*100.0:05.2f}%"
            )

            # 固定 10 epoch
            # if len(val_scores) == 1 or val_acc > val_scores[best_val_epoch]:
            #     print("\t   (New best performance, saving model...)")
            #     save_model(net, CHECKPOINT_PATH, model_name)
            #     best_val_epoch = epoch
            # elif best_val_epoch <= epoch - patience:
            #     print(f"Early stopping due to no improvement over the last {patience} epochs")
            #     break

        # Plot a curve of the validation accuracy
        plt.plot([i for i in range(1, len(val_scores) + 1)], val_scores)
        plt.xlabel("Epochs")
        plt.ylabel("Validation accuracy")
        plt.title(f"Validation performance of {model_name}")
        plt.show()
        plt.close()

    # load_model(CHECKPOINT_PATH, model_name, net=net)
    test_acc = test_model(net, test_loader)
    print((f" Test accuracy: {test_acc*100.0:4.2f}% ").center(50, "=") + "\n")
    return test_acc


def test_model(net, data_loader):
    """Test a model on a specified dataset.

    Args:
        net: Trained model of type BaseNetwork
        data_loader: DataLoader object of the dataset to test on (validation or test)
    """
    net.eval()
    true_preds, count = 0.0, 0
    for imgs, labels in data_loader:
        imgs, labels = imgs.to(device), labels.to(device)
        with torch.no_grad():
            preds = net(imgs).argmax(dim=-1)
            true_preds += (preds == labels).sum().item()
            count += labels.shape[0]
    test_acc = true_preds / count
    return test_acc



def visualize_activations_after_train(net, color="C0"):
    activations = {}

    net.eval()
    small_loader = data.DataLoader(train_set, batch_size=1024)
    imgs, labels = next(iter(small_loader))
    with torch.no_grad():
        layer_index = 0
        imgs = imgs.to(device)
        imgs = imgs.view(imgs.size(0), -1)
        # We need to manually loop through the layers to save all activations
        for layer_index, layer in enumerate(net.layers[:-1]):
            imgs = layer(imgs)
            activations[layer_index] = imgs.view(-1).cpu().numpy()

    # Plotting
    columns = 4
    rows = math.ceil(len(activations) / columns)
    fig, ax = plt.subplots(rows, columns, figsize=(columns * 2.7, rows * 2.5))
    fig_index = 0
    for key in activations:
        key_ax = ax[fig_index // columns][fig_index % columns]
        sns.histplot(data=activations[key], bins=50, ax=key_ax, color=color, kde=True, stat="density")
        key_ax.set_title(f"Layer {key} - {net.layers[key].__class__.__name__}")
        fig_index += 1
    fig.suptitle(f"Activation distribution for activation function {net.config['act_fn']['name']}", fontsize=14)
    fig.subplots_adjust(hspace=0.4, wspace=0.4)
    plt.show()
    plt.close()


def visualize_activations_before_train(act_fn_by_name):
    for i, act_fn_name in enumerate(act_fn_by_name):
        # Setting the seed ensures that we have the same weight initialization for each activation function
        act_fn = act_fn_by_name[act_fn_name]()
        net_actfn = BaseNetwork(act_fn=act_fn).to(device)
        visualize_gradients(net_actfn, color=f"C{i}")
def main():

    act_fn_by_name = {"sigmoid": Sigmoid,
                      "tanh": Tanh,
                      "relu": ReLU,
                      "leakyrelu": LeakyReLU,
                      "elu": ELU,
                      "SiLU": SiLU,
                      "selu": SELU,
                      "gelu": GELU,
                      "newgelu": NewGELU,
                      "fastgelu": FastGELU,
                      "quickgelu": QuickGELU,
                      "Mish": Mish,
                      }
    # visualize_activations_before_train(act_fn_by_name)
    # plot_all_actfn(act_fn_by_name)
    # plot_fashionmnist_example()

    for act_fn_name in act_fn_by_name:
        print(f"Training BaseNetwork with {act_fn_name} activation...")
        set_seed(42)
        act_fn = act_fn_by_name[act_fn_name]()
        net_actfn = BaseNetwork(act_fn=act_fn).to(device)
        train_model(net_actfn, f"FashionMNIST_{act_fn_name}", overwrite=False)

    # visualize_activations_after_train()

if __name__ == '__main__':
    set_seed(42)
    device = torch.device("cpu") if not torch.cuda.is_available() else torch.device("cuda:0")
    print("Using device", device)
    warnings.filterwarnings("ignore")


    transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])
    DATASET_PATH = 'fashionmnistdata'
    # Loading the training dataset. We need to split it into a training and validation part
    train_dataset = FashionMNIST(root=DATASET_PATH, train=True, transform=transform, download=True)
    train_set, val_set = torch.utils.data.random_split(train_dataset, [50000, 10000])
    # Loading the test set
    test_set = FashionMNIST(root=DATASET_PATH, train=False, transform=transform, download=True)

    train_loader = data.DataLoader(train_set, batch_size=1024, shuffle=True, drop_last=False)
    val_loader = data.DataLoader(val_set, batch_size=1024, shuffle=False, drop_last=False)
    test_loader = data.DataLoader(test_set, batch_size=1024, shuffle=False, drop_last=False)

    main()
